import pandas as pd
from pymilvus import connections, Collection, utility
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

# --- 配置参数 ---

# Milvus 连接配置
MILVUS_URI = "http://111.229.125.24:19530"
MILVUS_DB_NAME = "customer_service"
MILVUS_TOKEN = "root:Milvus"
COLLECTION_NAME = "general_faq"

# Excel 文件路径
EXCEL_FILE_PATH = "import_data_template.xlsx"

# 用于生成向量的预训练模型
# 'paraphrase-multilingual-MiniLM-L12-v2' 是一个优秀的多语言模型，支持中文
# 'BAAI/bge-large-zh-v1.5' 是一个顶尖的中文模型，维度为1024，如果使用需要修改Collection Schema
# 注意：首次运行时，会自动从网络下载模型，可能需要一些时间。
QUESTION_MODEL_NAME = 'paraphrase-multilingual-MiniLM-L12-v2'  # 输出384维向量
KEYWORDS_MODEL_NAME = 'paraphrase-multilingual-MiniLM-L12-v2'  # 输出384维向量


# --- 主逻辑 ---

def get_models():
    """加载并返回SentenceTransformer模型"""
    print("正在加载向量模型...")
    # 可以将两个模型设置为同一个，以节省内存
    question_model = SentenceTransformer(QUESTION_MODEL_NAME)
    keywords_model = SentenceTransformer(KEYWORDS_MODEL_NAME)
    print("模型加载完成。")
    return question_model, keywords_model

def connect_to_milvus():
    """连接到Milvus服务器"""
    print(f"正在连接到 Milvus (URI: {MILVUS_URI})...")
    connections.connect(
        uri=MILVUS_URI,
        db_name=MILVUS_DB_NAME,
        token=MILVUS_TOKEN
    )
    print("成功连接到 Milvus。")
    # 检查Collection是否存在
    if not utility.has_collection(COLLECTION_NAME):
        print(f"错误：集合 '{COLLECTION_NAME}' 在数据库 '{MILVUS_DB_NAME}' 中不存在。")
        print("请先运行 'create_collection.py' 脚本创建集合。")
        return None
    
    collection = Collection(name=COLLECTION_NAME)
    return collection

def prepare_data_for_insertion(df, question_model, keywords_model):
    """
    准备用于插入Milvus的数据列表。
    为 'question' 和 'keywords' 列生成向量。
    """
    print("正在准备数据并生成向量...")
    
    # 从DataFrame中提取文本数据
    questions = df['question'].tolist()
    # 将关键词列表合并为字符串进行编码
    keywords_texts = df['keywords'].apply(lambda x: ' '.join(str(x).split(','))).tolist()

    # 批量生成向量，使用tqdm显示进度条
    print("正在为 'question' 列生成向量...")
    question_embeddings = question_model.encode(questions, show_progress_bar=True)
    
    print("正在为 'keywords' 列生成向量...")
    keywords_embeddings = keywords_model.encode(keywords_texts, show_progress_bar=True)
    
    # 将处理好的数据整合到列表中
    data_to_insert = []
    for i in range(len(df)):
        entity = {
            "faq_id": df['faq_id'][i],
            "question_embedding": question_embeddings[i],
            "keywords_embedding": keywords_embeddings[i],
            "question": df['question'][i],
            "answer": df['answer'][i],
            # 将逗号分隔的字符串转为列表
            "keywords": str(df['keywords'][i]).split(','),
            "category": df['category'][i],
            "scene_tags": str(df['scene_tags'][i]).split(','),
            "emotion_type": df['emotion_type'][i],
            "priority": df['priority'][i],
            "use_count": df['use_count'][i],
            "satisfaction_rate": df['satisfaction_rate'][i],
            "update_time": df['update_time'][i]
        }
        data_to_insert.append(entity)
        
    print(f"数据准备完成，共 {len(data_to_insert)} 条。")
    return data_to_insert

def main():
    """
    主函数：读取Excel，生成向量，并导入Milvus
    """
    # 1. 加载模型
    question_model, keywords_model = get_models()

    # 2. 连接到Milvus并获取Collection对象
    collection = connect_to_milvus()
    if not collection:
        return

    # 3. 从Excel读取数据
    try:
        print(f"正在从 '{EXCEL_FILE_PATH}' 读取数据...")
        df = pd.read_excel(EXCEL_FILE_PATH)
        # 基础的数据校验
        if 'faq_id' not in df.columns or 'question' not in df.columns:
            print("错误：Excel文件中缺少'faq_id'或'question'列。")
            return
        print(f"成功读取 {len(df)} 条数据。")
    except FileNotFoundError:
        print(f"错误：找不到文件 '{EXCEL_FILE_PATH}'。请确保文件存在于当前目录。")
        return

    # 4. 准备数据
    data_list = prepare_data_for_insertion(df, question_model, keywords_model)

    # 5. 插入数据到Milvus
    print("正在向 Milvus 集合中插入数据...")
    try:
        # Milvus的insert方法支持直接传入字典列表
        result = collection.insert(data_list)
        print("数据插入成功！")
        print(f"主键(Primary Keys): {result.primary_keys}")
        print(f"成功插入行数: {result.insert_count}")
    except Exception as e:
        print(f"数据插入时发生错误: {e}")

    # 6. 断开连接
    connections.disconnect("default")
    print("与 Milvus 的连接已关闭。")

if __name__ == "__main__":
    main() 